#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/platform/logging.h"

#include <cstring>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <vector>
#include <string>
#include <iostream>
#include <algorithm>

#include "slot_line_util.h"

using namespace tensorflow;
namespace tensorflow {

REGISTER_OP("ParseSampleV1")
    .Attr("hash_slots: list(int)")
    .Attr("hash_slots_size: int")
    .Attr("key_max_size: int")
    .Input("line: string")
    .Output("slot_value: hash_slots_size * int64")
    .Output("slot_indices: hash_slots_size * int64")
    .Output("slot_shape: hash_slots_size * int64")
    .Output("label: float32")
    .Output("key: string")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
        shape_inference::ShapeHandle input_shape;
        TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input_shape));
        shape_inference::DimensionHandle row_shape = c->Dim(input_shape, 0);
        int32 hash_slots_size;
        TF_RETURN_IF_ERROR(c->GetAttr("hash_slots_size", &hash_slots_size));
        int output_index = 0;
        // slot values
        for (int i = 0; i < hash_slots_size; i++) {
            c->set_output(output_index++, c->Vector(c->UnknownDim()));
        }

        // slot indices
        for (int i = 0; i < hash_slots_size; i++) {
            c->set_output(output_index++, c->Matrix(c->UnknownDim(), 2));
        }

        // slot shape
        for (int i = 0; i < hash_slots_size; i++) {
            c->set_output(output_index++, c->Vector(2));
        }

        // set label
        c->set_output(output_index++, c->Matrix(row_shape, c->UnknownDim()));

        // key
        c->set_output(output_index++, c->Vector(c->UnknownDim()));

        return Status::OK();
    });

class ParseSampleV1 : public OpKernel {
private:
    vector<int> hash_slots;
    int64 key_max_size;
public:
    explicit ParseSampleV1(OpKernelConstruction* context) : OpKernel(context) {
        OP_REQUIRES_OK(context, context->GetAttr("hash_slots", &hash_slots));
        OP_REQUIRES_OK(context, context->GetAttr("key_max_size", &key_max_size));
    }

    void Compute(OpKernelContext* context) override {
        // Grab the input tensor
        const Tensor& input_tensor = context->input(0);
        auto input = input_tensor.flat<string>();
        auto input_size = input.size();

        // parse lines
        unordered_map<int, vector<int64_t> > hash_slot_values;
        unordered_map<int, vector<pair<int, int> > > hash_slot_indices;
        unordered_map<int, int> hash_slot_shapes;
        vector<vector<float> > labels;
        vector<string> keys;

        for (int i = 0; i < input.size(); i++) {
            string line = input(i);
            vector<float> label;
            string key;
            ParseBatchSparse(line, i, label, 
                            hash_slot_values, hash_slot_indices, hash_slot_shapes,
                            key);
            keys.push_back(move(key));
            labels.push_back(move(label));
        }

        int output_index = 0;
        // set values
        TensorShape empty_shape({1});
        TensorShape empty_indices({1, 2});
        TensorShape slot_shape({2});

        int slot_value_output_offset = 0;
        int slot_indices_output_offset = this->hash_slots.size();
        int slot_shape_output_offset = this->hash_slots.size()*2;

        for (int i = 0; i < this->hash_slots.size(); i++) {
            const auto& slot_id = this->hash_slots[i];
            const auto& value_iter = hash_slot_values.find(slot_id);
            const auto& indices_iter = hash_slot_indices.find(slot_id);
            const auto& slot_shape_iter = hash_slot_shapes.find(slot_id);

            if (value_iter == hash_slot_values.end()) {
                Tensor* value_tensor = NULL;
                Tensor* indices_tensor = NULL;
                Tensor* shape_tensor = NULL;
                OP_REQUIRES_OK(context, context->allocate_output(i+slot_value_output_offset, empty_shape, &value_tensor));
                output_index++;
                OP_REQUIRES_OK(context, context->allocate_output(i+slot_indices_output_offset, empty_indices, &indices_tensor));
                output_index++;
                OP_REQUIRES_OK(context, context->allocate_output(i+slot_shape_output_offset, slot_shape, &shape_tensor));
                output_index++;
                value_tensor->flat<int64>()(0) = -1;
                indices_tensor->flat<int64>()(0) = 0;
                indices_tensor->flat<int64>()(1) = 0;
                SetShapeTensor(shape_tensor, input_size, 1);
            } else {
                long value_size = value_iter->second.size();
                TensorShape value_shape({value_size});
                TensorShape indices_shape({value_size, 2});
                Tensor* value_tensor = NULL;
                Tensor* indices_tensor = NULL;
                Tensor* shape_tensor = NULL;
                OP_REQUIRES_OK(context, context->allocate_output(i+slot_value_output_offset, value_shape, &value_tensor));
                output_index++;
                OP_REQUIRES_OK(context, context->allocate_output(i+slot_indices_output_offset, indices_shape, &indices_tensor));
                output_index++;
                OP_REQUIRES_OK(context, context->allocate_output(i+slot_shape_output_offset, slot_shape, &shape_tensor));
                output_index++;
                auto value_flat = value_tensor->flat<int64>();
                auto indices_mat = indices_tensor->tensor<int64, 2>();
                // set indices mat
                for (int i = 0; i < value_size; i++) {
                    value_flat(i) = value_iter->second[i];
                    auto const& pair = indices_iter->second[i];
                    indices_mat(i, 0) = pair.first;
                    indices_mat(i, 1) = pair.second;
                }
                SetShapeTensor(shape_tensor, input_size, slot_shape_iter->second);
            }

        }

        // set label
        long label_size = labels[0].size();
        TensorShape label_shape({input.size(), label_size});
        Tensor* label_tensor = NULL;
        OP_REQUIRES_OK(context, context->allocate_output(output_index, label_shape, &label_tensor));

        auto label_output = label_tensor->tensor<float, 2>();
        for (int i = 0; i < input_size; i++) {
            for (int j = 0; j < label_size; j++) {
                label_output(i, j) = labels[i][j];
            }
        }
        output_index++;

        // set key
        TensorShape key_shape({input.size()});
        Tensor* key_tensor = NULL;
        OP_REQUIRES_OK(context, context->allocate_output(output_index, key_shape, &key_tensor));
        auto key_tensor_flatten = key_tensor->flat<string>();
        for (int i = 0; i < input_size; i++) {
            key_tensor_flatten(i) = move(keys[i]);
        }
        output_index++;

    }

    void SetShapeTensor(Tensor* shape_tensor,
                        const int& input_size,
                        const int& column) {
        auto shape_out = shape_tensor->flat<int64>();
        shape_out(0) = input_size;
        shape_out(1) = column;

    }

    bool InHashSlots(int slotid) {
        vector<int>::iterator it = find(hash_slots.begin(), hash_slots.end(), slotid);
        if (it != hash_slots.end()) {
            return true;
        } else {
            return false;
        }
    }

    void ParseBatchSparse(const string& slot_line,
                        const int& batch_index,
                        vector<float>& labels,
                        unordered_map<int, vector<int64_t> >& hash_slot_values,
                        unordered_map<int, vector<pair<int, int> > >& hash_slot_indices,
                        unordered_map<int, int>& hash_slot_shapes,
                        string& key
                        ) {
        vector<string> slot_tokens;
        tf_op::SlotLineSplit(slot_line, slot_tokens, "\u0001");
        // sample key
        key = move(slot_tokens[0]);
        // add labels
        const auto& label_token = slot_tokens[2];
        vector<string> label_tokens;
        tf_op::SlotLineSplit(label_token, label_tokens, ",");
        for (const auto& label_value : label_tokens) {
            labels.push_back(atof(label_value.c_str()));
        }

        // add slot features
        unordered_map<int, int> current_slot_indices;
        vector<string> feas;
        tf_op::SlotLineSplit(slot_tokens[3], feas, "\u0002");
        for (int i = 0; i < feas.size(); i++) {
            vector<string> v_tokens;
            tf_op::SlotLineSplit(feas[i], v_tokens, "\u0003");
            int slot_id = strtol(v_tokens[0].c_str(), nullptr, 10);
            if (InHashSlots(slot_id)) {
                if (hash_slot_values.count(slot_id)==0) {
                    vector<int64_t> values;
                    hash_slot_values[slot_id] = values;
                }

                if (hash_slot_indices.count(slot_id)==0) {
                    vector<pair<int, int> > indices;
                    hash_slot_indices[slot_id] = indices;
                }

                if (hash_slot_shapes.count(slot_id)==0) {
                    hash_slot_shapes[slot_id] = 0;
                }

                if (current_slot_indices.count(slot_id)==0) {
                    current_slot_indices[slot_id] = 0;
                }

                // fe那边hash是uint64
                // hash_slot_values[slot_id].push_back(strtol(v_tokens[1].c_str(), nullptr, 10));
                int64_t old_key = int64_t(strtoul(v_tokens[1].c_str(), nullptr, 10) & 9223372036854775807);
                int64_t mod_key = old_key % this->key_max_size;
                hash_slot_values[slot_id].push_back(mod_key);
                auto& indice_size = current_slot_indices[slot_id];
                hash_slot_indices[slot_id].push_back(std::make_pair(batch_index, indice_size));
                indice_size++;
                if (hash_slot_shapes[slot_id] < indice_size) {
                    hash_slot_shapes[slot_id] = indice_size;
                }

            } else {
                continue;
            }
        }

    }
};



REGISTER_KERNEL_BUILDER(Name("ParseSampleV1").Device(DEVICE_CPU), ParseSampleV1);


}









